package cn.itwen.star.graph.core.service.impl;

import cn.hutool.core.date.DateTime;
import cn.itwen.star.graph.comfyui.client.pojo.ComfyuiTask;
import cn.itwen.star.graph.comfyui.client.pojo.MessageBase;
import cn.itwen.star.graph.core.job.RunTaskJob;
import cn.itwen.star.graph.core.service.ComfyuiMessageService;
import cn.itwen.star.graph.core.service.RedisService;
import cn.itwen.star.graph.core.service.UserResultService;
import cn.itwen.star.graph.core.service.WsNoticeService;
import com.alibaba.fastjson2.JSON;
import org.redisson.api.RSemaphore;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;

@Service
public class ComfyuiMessageServiceImpl implements ComfyuiMessageService {

    @Autowired
    private RedisService redisService;

    @Autowired
    private WsNoticeService wsNoticeService;

    @Autowired
    private UserResultService userResultService;

    @Autowired
    private RedissonClient redissonClient;

    @Override
    public void handleMessage(MessageBase messageBase) {
        if ("progress".equals(messageBase.getType())){
            //处理文生图片进度
            handleProgress(messageBase);
        } else if ("executed".equals(messageBase.getType())) {
            //处理文生图片结果
            handleExecuted(messageBase);
        } else if ("execution_error".equals(messageBase.getType())) {
            //处理文生图片错误
            handleError(messageBase);
        } else if ("status".equals(messageBase.getType())) {
            //处理文生图片状态
            handleStatus(messageBase);

        }

    }

    /**
     * 处理文生图片状态
     * @param messageBase
     */
    private void handleStatus(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData();
        HashMap<String, Object> status = (HashMap<String, Object>) data.get("status");
        HashMap<String, Object> execInfo = (HashMap<String, Object>)status.get("exec_info");
        Integer queueRemaining = (Integer) execInfo.get("queue_remaining");

        //如果队列的数量为0，则释放信号量
        if (queueRemaining == 0){
            RSemaphore semaphore = redissonClient.getSemaphore(RunTaskJob.TASK_RUN_SEMAPHORE);
            semaphore.release();
            System.out.println(DateTime.now() + "\t\t释放许可证数量：" + semaphore.availablePermits());
        }
    }

    /**
     * 处理文生图片错误
     * @param messageBase
     */
    private void handleError(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData();
        String promptId = data.get("prompt_id").toString();
        wsNoticeService.sendError(promptId,data);
    }

    /**
     * 处理文生图片结果
     * @param messageBase
     */
    private void handleExecuted(MessageBase messageBase) {
        //1.消息格式的转换
        HashMap<String, Object> data = messageBase.getData();
        HashMap<String, Object> output = (HashMap<String, Object>) data.get("output");
        List<HashMap<String, Object>> images = (List<HashMap<String, Object>>) output.get("images");
        List<String> urls = images.stream().map(image ->
                String.format("http://127.0.0.1:8188/view?filename=%s&type=%s&subfloder=%s", image.get("filename"), image.get("type"), image.get("subfolder"))
        ).toList();
        //2.从redis中获取任务
        String promptId = data.get("prompt_id").toString();
        ComfyuiTask task = redisService.getStartedTask(promptId);
        if (task == null){
            return;
        }
        //3.推送消息给指定用户
        HashMap<String, Object> result = new HashMap<>();
        result.put("type", "imageResult");
        result.put("urls",urls);
        wsNoticeService.sendToUser(task.getWsClientId(), JSON.toJSONString(result));

        //4.调用UserResultService保存结果
        userResultService.insertUserResult(task,urls);
    }

    /**
     * 处理文生图片进度
     * @param messageBase
     */
    private void handleProgress(MessageBase messageBase) {
        //1.消息格式转换
        HashMap<String, Object> data = messageBase.getData();
        data.put("type","progress");

        //2.从redis中获取任务（客户端id）
        String promptId = (String) data.get("prompt_id");
        ComfyuiTask task = redisService.getStartedTask(promptId);
        if (task ==  null){
            return;
        }
        //3.发送给客户端
        wsNoticeService.sendToUser(task.getWsClientId(), JSON.toJSONString(data));
    }
}
